from flask_sqlalchemy import SQLAlchemy
from sqlalchemy import Column, String, Integer, ForeignKey, and_
from sqlalchemy.types import BINARY
from database import db
from os import path, mkdir
from config import storage_path
import re, os
import Encryption_and_decryption as ed
from urllib.parse import quote

filename_pattern = re.compile(r"[^\u4e00-\u9fa5]+")

# 定义允许的文件后缀列表
allowed_file_suffix_list = {
    "doc",
    "docx",
    "ppt",
    "pptx",
    "xls",
    "xlsx",
    "pdf",
    "png",
    "jpg",
    "jpeg",
    "gif",
}


# 文件资源表
class File(db.Model):
    __tablename__ = "file"
    file_id = db.Column(db.Integer, primary_key=True)
    user_id = db.Column(
        db.Integer, db.ForeignKey("user.user_id", ondelete="CASCADE"), nullable=False
    )
    file_name = db.Column(db.String(256))

    # 对密文的HMAC数字签名
    file_signature = db.Column(db.BINARY(128))
    # 明文 静态散列值
    static_hash = db.Column(db.BINARY(128))
    # 验证文件完整性公钥
    public_key = db.Column(db.BINARY(32))
    

    # 类的内置函数
    @classmethod
    def upload_file(cls, user, data):
        # 文件名验证操作：
        filename = data.filename

        # 检查文件的后缀是否在允许的文件后缀列表中
        filename_suffix = filename.rsplit(".", maxsplit=1)[-1]
        assert filename_suffix in allowed_file_suffix_list, "banned file type"

        # 读取文件的内容
        content = data.read()

        # 验证文件的大小是否小于 10MB
        assert len(content) < 10 * 1024 * 1024, "file too large (>=10MB)"

        # 创建单个用户的文件夹路径
        user_id = str(user.user_id) + "/"
        if not path.exists(storage_path + user_id):
            if not path.exists(storage_path):
                mkdir(storage_path)
            mkdir(storage_path + user_id)

        # 用户id
        user_id_ = user.user_id
        # 生成盐值
        user_salt = user.user_salt
        # 密钥哈希值
        hash_password_ = user.hash_password
        # 利用对称加密生成密文
        file_info_ = ed.file_upload_encryption(hash_password_, content)
        # 静态散列值
        static_hash = ed.get_hash_from_bytes(content)
        # 对密文的签名
        file_signature_ = ed.Hash_get(hash_password_, user_salt, file_info_)[1]


        # 验证文件完整性公钥
        public_key_ = ed.Hash_get(hash_password_, user_salt, file_info_)[0]

        # 文件上传：
        # 将加密后的文件保存到用户文件夹中
        file_path = os.path.join(storage_path + str(user_id_), filename)
        with open(file_path, "wb") as f:
            f.write(file_info_)

        file = File(
            user_id=user_id_,
            file_name=filename,
            file_signature=file_signature_,
            static_hash=static_hash,
            public_key=public_key_
        )
        db.session.add(file)
        # 提交数据库更改
        db.session.commit()

    @classmethod
    def download_file(cls, user, filename, type):
        from flask import make_response

        # 获取对应的文件
        find_file = File.query.filter(
            and_(File.user_id == user.user_id, File.file_name == filename)
        ).first()

        # 读取加密后的文件内容
        file_path = os.path.join(storage_path + str(user.user_id), filename)
        with open(file_path, "rb") as f:
            file_content = f.read()

        # 加密文件：
        encrypted_file = file_content

        # 密钥哈希值
        hash_password_ = user.hash_password

        # 完整性验证部分
        if ed.Hash_verify(find_file.public_key,find_file.file_signature,encrypted_file):
            # 利用对称密钥对密文解密,获取明文
            decrypted_file = ed.file_download_decryption(hash_password_, file_content)
        else:
            return '完整性验证失败！'

        # 判断要下载的类型
        if type == "hashvalue":
            content = find_file.static_hash
            filename = "hash_" + filename

            # 创建content 响应对象
            response = make_response(content)
            # 去掉文件名本身的后缀
            file_basename = os.path.splitext(filename)[0]
            # 对文件名进行URL编码，并设置为附件的下载文件名
            url_filename = quote(file_basename) + ".txt"
            response.headers[
                "Content-Disposition"
            ] = f'attachment; filename="{url_filename}"'
            # 返回响应对象，完成下载功能
            return response

        elif type == "signature":
            # 读取签名
            content = find_file.file_signature
            filename = "sig_" + filename

            # 创建content 响应对象
            response = make_response(content)
            # 去掉文件名本身的后缀
            file_basename = os.path.splitext(filename)[0]
            # 对文件名进行URL编码，并设置为附件的下载文件名
            url_filename = quote(file_basename) + ".txt"
            response.headers[
                "Content-Disposition"
            ] = f'attachment; filename="{url_filename}"'
            # 返回响应对象，完成下载功能
            return response

        else:
            if type == "plaintext":
                # 解密并下载
                content = decrypted_file
                filename = "decrypted_" + filename
                
                # 创建content 响应对象
                response = make_response(content)
                # 对文件名进行URL编码，并设置为附件的下载文件名
                url_filename = quote(filename)
                response.headers['Content-Disposition'] = f'attachment; filename="{url_filename}"'
                # 返回响应对象，完成下载功能
                return response


            elif type == "encrypted":
                content = encrypted_file
                filename = "encrypted_" + filename

                # 创建content 响应对象
                response = make_response(content)
                # 对文件名进行URL编码，并设置为附件的下载文件名
                url_filename = quote(filename)
                response.headers[
                    "Content-Disposition"
                ] = f'attachment; filename="{url_filename}"'
                # 返回响应对象，完成下载功能
                return response
